#include "Ampere_Q80.h"

static AMPERE_FLASH_MEMMAP Ampflash = {0};
static int CpuNum = 0;

static AMPERE_FLASH_MEMMAP *get_VirtAddr_from_context(const struct flashctx *flash)
{
  if (!flash || !flash->mst || !flash->mst->spi.data)
  {
    msg_perr("Unable to extract Virt Addr from flash context.\n");
    return NULL;
  }
  return flash->mst->spi.data;
}

static int Amp_spi_send_command(const struct flashctx *flash,
                                unsigned int writecnt, unsigned int readcnt,
                                const unsigned char *writearr,
                                unsigned char *readarr)
{
  uint32_t cmd;
  uint32_t Addr;
  uint32_t read[2];
  uint32_t write[2];
  uint32_t Num = 0;

  if (writecnt > 9 || readcnt > 8 || writecnt == 0)
  {
    msg_perr("%s: Invalid read/write count for send command.\n",
             __func__);
    return SPI_GENERIC_ERROR;
  }

  AMPERE_FLASH_MEMMAP *pVirtAddr = get_VirtAddr_from_context(flash);
  if (NULL == pVirtAddr)
    return SPI_GENERIC_ERROR;

  if ((ERASE_4 == writearr[0]) || (ERASE_32 == writearr[0]) || (ERASE_64 == writearr[0]) || (ERASE_D7 == writearr[0]) ||
      (ERASE_C4 == writearr[0]) || (ERASE_DB == writearr[0]) || (ERASE_50 == writearr[0]) || (ERASE_81 == writearr[0]) ||
      (ERASE_21 == writearr[0]) || (ERASE_5C == writearr[0]) || (ERASE_DC == writearr[0]))

  {
    if (writecnt < 4)
    {
      msg_perr("%s: Address Info Not exit.\n", __func__);
      return SPI_GENERIC_ERROR;
    }

    if (4 == writecnt)
    {
      cmd = ((writearr[0] << 8) | (0x1 << 6) | (0x0 << 5));
      MmioWrite32(pVirtAddr->spi_ctrl_reg_base_virt + SPI_CMD_PORT_ADDR, cmd);
      Addr = ((writearr[1] << 16) | (writearr[2] << 8) | (writearr[3]));
    }
    else
    {
      cmd = ((writearr[0] << 8) | (0x1 << 6) | (0x1 << 5));
      MmioWrite32(pVirtAddr->spi_ctrl_reg_base_virt + SPI_CMD_PORT_ADDR, cmd);
      Addr = ((writearr[1] << 16) | (writearr[2] << 16) | (writearr[3] << 8) | (writearr[4]));
    }

    // msg_perr("Addr  %x %x.\n", Addr, writecnt);

    MmioWrite32(pVirtAddr->spi_ctrl_reg_base_virt + SPI_ADDR_PORT_ADDR, Addr);
    // asm volatile ("isb sy":::"cc");
    // asm volatile ("dsb sy":::"cc");

    MmioWrite32(pVirtAddr->spi_ctrl_reg_base_virt + SPI_LDATA_PORT_ADDR, 0x1);
    // asm volatile ("isb sy":::"cc");
    // asm volatile ("dsb sy":::"cc");
    return 0;
  }

  if (1 == writecnt)
  {
    cmd = (writearr[0] << 8);
    MmioWrite32(pVirtAddr->spi_ctrl_reg_base_virt + SPI_CMD_PORT_ADDR, cmd);
    // msg_perr("cmd %x.\n", cmd);
    // asm volatile ("isb sy":::"cc");
    // asm volatile ("dsb sy":::"cc");
    MmioWrite32(pVirtAddr->spi_ctrl_reg_base_virt + SPI_LDATA_PORT_ADDR, 0x1);
    // asm volatile ("isb sy":::"cc");
    // asm volatile ("dsb sy":::"cc");
  }
  else if (writecnt < 6)
  {
    memcpy(&write[0], &writearr[1], writecnt - 1);
  }
  else
  {
    memcpy(&write[1], &writearr[5], writecnt - 5);
    memcpy(&write[0], &writearr[1], 4);
  }

  if (writecnt > 1)
  {
    Num = (writecnt - 1);
    cmd = ((writearr[0] << 8) | Num);
    MmioWrite32(pVirtAddr->spi_ctrl_reg_base_virt + SPI_CMD_PORT_ADDR, cmd);
    // asm volatile ("isb sy":::"cc");
    // asm volatile ("dsb sy":::"cc");
    MmioWrite32(pVirtAddr->spi_ctrl_reg_base_virt + SPI_HDATA_PORT_ADDR, write[1]);
    // asm volatile ("isb sy":::"cc");
    // asm volatile ("dsb sy":::"cc");
    MmioWrite32(pVirtAddr->spi_ctrl_reg_base_virt + SPI_LDATA_PORT_ADDR, write[0]);
    // asm volatile ("isb sy":::"cc");
    // asm volatile ("dsb sy":::"cc");
  }

  // msg_perr("write %x %d.\n", writearr[1], writecnt);

  if (readcnt)
  {
    Num = ((readcnt > 2) ? ((readcnt + 4 - 1) / 4 * 4) : readcnt);
    cmd = ((writearr[0] << 8) | Num);
    MmioWrite32(pVirtAddr->spi_ctrl_reg_base_virt + SPI_CMD_PORT_ADDR, cmd);
    // msg_perr("cmd %x.\n", cmd);
    // asm volatile ("isb sy":::"cc");
    // asm volatile ("dsb sy":::"cc");
    read[0] = MmioRead32(pVirtAddr->spi_ctrl_reg_base_virt + SPI_LDATA_PORT_ADDR);
    read[1] = MmioRead32(pVirtAddr->spi_ctrl_reg_base_virt + SPI_HDATA_PORT_ADDR);
    // asm volatile ("isb sy":::"cc");
    // asm volatile ("dsb sy":::"cc");
    MmioWrite32(pVirtAddr->spi_ctrl_reg_base_virt + SPI_LDATA_PORT_ADDR, 0x1);
    // asm volatile ("isb sy":::"cc");
    // asm volatile ("dsb sy":::"cc");
    // msg_perr("actu read %8x %8x.\n", read[0], read[1]);
  }

  if (0 == readcnt)
  {
    return 0;
  }
  else if (readcnt < 5)
  {
    memcpy(readarr, &read[0], readcnt);
  }
  else
  {
    memcpy((readarr + 4), &read[1], readcnt - 4);
    memcpy(readarr, &read[0], 4);
  }

  return 0;
}

static void spi_pp_cfg(struct flashctx *flash)
{
  AMPERE_FLASH_MEMMAP *pVirtAddr = get_VirtAddr_from_context(flash);

  int value = 0;
  value &= (~(SPI_ACCESS_CMD_MODE_MASK | SPI_ACCESS_ADDR_MODE_MASK));
  value = (0x0 << 6) | (0x0 << 5) | (0x0 << 3); /* pp(0x2) */

  MmioWrite32(pVirtAddr->spi_ctrl_reg_base_virt + SPI_CFG_PORT_ADDR, value);
  MmioWrite32(pVirtAddr->spi_ctrl_reg_base_virt + SPI_FLASH_CAPACITY, 0x00);
  pVirtAddr->size = (flash->chip->total_size * 1024);
  if (NULL == pVirtAddr->spi_rd_base_virt)
  {
    pVirtAddr->spi_rd_base_virt = physmap("QSPIRead", (SPI_READ_BASE + CpuNum * SPI_CPU_INTERVAL), pVirtAddr->size);
  }

  if (NULL == pVirtAddr->spi_wt_base_virt)
  {
    pVirtAddr->spi_wt_base_virt = physmap("QSPIWrite", (SPI_WRITE_BASE + CpuNum * SPI_CPU_INTERVAL), pVirtAddr->size);
  }

  return;
}

static int Amp_spi_read(struct flashctx *flash, uint8_t *buf,
                        unsigned int start, unsigned int len)
{
  AMPERE_FLASH_MEMMAP *pVirtAddr = get_VirtAddr_from_context(flash);
  if (NULL == pVirtAddr)
    return SPI_GENERIC_ERROR;

  spi_pp_cfg(flash);

  if (start > pVirtAddr->size)
  {
    memset(buf, 0xFF, len);
    return 0;
  }

  if ((len & 0x07) || (start & 0x07))
  {
    //
    // When count is not multiple of 8 bytes, memcpy() function may cause bus error.
    uint8_t *ptr;
    uint32_t indx;

    ptr = (uint8_t *)(pVirtAddr->spi_rd_base_virt + start);
    for (indx = 0; indx < len; indx++)
      buf[indx] = ptr[indx];
  }
  else
  {
    memcpy(buf, pVirtAddr->spi_rd_base_virt + start, len);
  }

  return 0;
}

static void AmpereSPI_Write_Word(struct flashctx *flash, UINTN Address, UINT32 Value)
{

  AMPERE_FLASH_MEMMAP *pVirtAddr = get_VirtAddr_from_context(flash);
  if (Address % 4 != 0)
  {
    msg_perr("Address not aligne 4 byte \n");
  }

  MmioWrite32(pVirtAddr->spi_ctrl_reg_base_virt + SPI_CMD_PORT_ADDR, (0x06 << 8));
  MmioWrite32(pVirtAddr->spi_ctrl_reg_base_virt + SPI_LDATA_PORT_ADDR, 0x1);
  MmioWrite32(Address, Value);
  // asm volatile ("isb sy":::"cc");
  // asm volatile ("dsb sy":::"cc");
  MmioWrite32(pVirtAddr->spi_ctrl_reg_base_virt + SPI_CMD_PORT_ADDR, (0x05 << 8 | 0x1));
  MmioWrite32(pVirtAddr->spi_ctrl_reg_base_virt + SPI_LDATA_PORT_ADDR, 0x1);
  while ((MmioRead32(pVirtAddr->spi_ctrl_reg_base_virt + SPI_LDATA_PORT_ADDR) & 0x1) == 0x1)
    ;
  // asm volatile ("isb sy":::"cc");
  // asm volatile ("dsb sy":::"cc");
}

static int Amp_spi_write_256(struct flashctx *flash, const uint8_t *buf,
                             unsigned int start, unsigned int len)
{
  uint32_t index;
  uintptr_t dest_addr;
  AMPERE_FLASH_MEMMAP *pVirtAddr = get_VirtAddr_from_context(flash);

  spi_pp_cfg(flash);
  if (start > pVirtAddr->size)
  {
    return 0;
  }

  dest_addr = (uintptr_t)pVirtAddr->spi_wt_base_virt + start;

  MmioWrite32(pVirtAddr->spi_ctrl_reg_base_virt + SPI_CMD_PORT_ADDR, (0x05 << 8 | 0x1));
  MmioWrite32(pVirtAddr->spi_ctrl_reg_base_virt + SPI_LDATA_PORT_ADDR, 0x1);
  while ((MmioRead32(pVirtAddr->spi_ctrl_reg_base_virt + SPI_LDATA_PORT_ADDR) & 0x1) == 0x1)
    ;

  for (index = 0; index < len; index += 4)
  {
    AmpereSPI_Write_Word(flash, dest_addr + index, *(uint32_t *)(buf + index));
  }

  MmioWrite32(pVirtAddr->spi_ctrl_reg_base_virt + SPI_FLASH_FLUSH, 0x1);

  return 0;
}

static int Amp_spi_write_aai(struct flashctx *flash, const uint8_t *buf,
                             unsigned int start, unsigned int len)
{
  msg_perr("%s: Not Support.\n", __func__);

  return 0;
}

static int Amp_spi_shutdown(void *data)
{
  AMPERE_FLASH_MEMMAP *pVirtAddr = data;
  if (pVirtAddr->spi_ctrl_reg_base_virt)
  {
    physunmap(pVirtAddr->spi_ctrl_reg_base_virt, SPI_REG_LENGTH);
    pVirtAddr->spi_ctrl_reg_base_virt = NULL;
  }

  if (pVirtAddr->spi_rd_base_virt)
  {
    physunmap(pVirtAddr->spi_rd_base_virt, pVirtAddr->size);
    pVirtAddr->spi_rd_base_virt = NULL;
  }

  if (pVirtAddr->spi_wt_base_virt)
  {
    physunmap(pVirtAddr->spi_wt_base_virt, pVirtAddr->size);
    pVirtAddr->spi_wt_base_virt = NULL;
  }

  return 0;
}

static const struct spi_master spi_master_Ampere = {
    .max_data_read = 9,
    .max_data_write = 9,
    .command = Amp_spi_send_command,
    .multicommand = default_spi_send_multicommand,
    .read = Amp_spi_read,
    .write_256 = Amp_spi_write_256,
    .write_aai = Amp_spi_write_aai,
    .shutdown = Amp_spi_shutdown,
};

static void *AmpMap(const char *descr, uintptr_t phys_addr, size_t len)
{
  return Ampflash.spi_rd_base_virt;
}

static void AmpunMap(void *virt_addr, size_t len)
{
  return;
}

static int Ampere_spi_init(void)
{
  char *arg;

  arg = extract_programmer_param("CPU");
  if (arg)
  {
    CpuNum = strtol(arg, NULL, 10);
  }
  else if (arg && !strlen(arg))
  {
    msg_perr("Missing argument for CPU.\n");
    free(arg);
    return 1;
  }

  if (arg)
    free(arg);

  Ampflash.spi_ctrl_reg_base_virt = physmap("QSPICtrl", (SPI_CTRL_REG_BASE + CpuNum * SPI_CPU_INTERVAL), SPI_REG_LENGTH);
  return register_spi_master(&spi_master_Ampere, &Ampflash);
}

const struct programmer_entry programmer_Ampere = {
    .name = "AMPERE_Q64",
    .type = OTHER,
    .devs.note = "AMPERE_Q64.\n",
    .init = Ampere_spi_init,
    .map_flash_region = AmpMap,
    .unmap_flash_region = AmpunMap,
    .delay = internal_delay,
};